'''
Created on May 21, 2012

@author: tel
'''
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from pylab import setp
import stats
import sys, copy, os, pickle
import scipy.stats as ss

def OffTicks(ax, list):
    for entry in list:
        if entry=='x':
            setp(ax.get_xticklabels(), visible=False)
        elif entry=='y':
            setp(ax.get_yticklabels(), visible=False)

def AddGamma(ax, data, rang):
#     histy, histx = np.histogram(data, range=rang)
    result = stats.fit_gamma(data, rang)
    textstr = (r'$\alpha=%.2f$'+'\n'+r'$\mathrm{scale}=%.2f$') % (result['shape'], result['scale'])
    
    x = np.arange(*rang)
    y = []
    for val in x:
        y.append(ss.gamma.pdf(val, result['shape'], loc=result['loc'], scale=result['scale'])*result['count'])
    ax.plot(x, y, 'g-.')
    
    # these are matplotlib.patch.Patch properies
    props = dict(boxstyle='round', facecolor='wheat', alpha=0.5)

    #place a text box in upper left in axes coords
    ax.text(0.05, 0.95, textstr, transform=ax.transAxes, fontsize=8,
            verticalalignment='top', bbox=props)

def TwoPanelHisto(l, r):
    fig = plt.figure()
    ax1 = fig.add_subplot(2,1,1)
    ax1.hist(l)
    ax2 = fig.add_subplot(2,1,2,sharex=ax1)
    ax2.hist(r)
    plt.show()

def UPSixPanelHisto(dir):
    f = open(os.path.join(dir, 'SixPanelHisto.pck'), 'r')
    unpickler = pickle.Unpickler(f)
    args = []
    for i in range(6):
        args.append(unpickler.load())
    f.close()
    SixPanelHisto(dir, *args, repickle=False) 

    
def Unpickle(indir, infile, outdir, outfile, plotter, argn):
    f = open(os.path.join(indir, infile+'.pck'),'r')
    Unpickler = pickle.Unpickler(f)
    args = []
    for i in range(argn):
        args.append(Unpickler.load())
    f.close()
    plotter(*args, repickle=False, dir=outdir, figname=outfile)

def Repickle(outdir, outfile, *args):
    f = open(os.path.join(outdir, outfile+'.pck'), 'w')
    pickler = pickle.Pickler(f)
    for list in args:
        pickler.dump(list)
    f.close()

def SixPanelHisto(ul, ur, ml, mr, ll, lr, repickle=True, outdir='home/tel/Dropbox/remote', outfile='6panelhisto'):
    lrange = (-15, 20)
    rrange = (-20, 15)
    bin_width = 1
    lbins = int((lrange[1]-lrange[0])/bin_width)
    rbins = int((rrange[1]-rrange[0])/bin_width)
    
    if repickle:
        Repickle(outdir, outfile, ul, ur, ml, mr, ll, lr)
#        f = open(os.path.join(dir, 'SixPanelHisto.pck'), 'w')
#        pickler = pickle.Pickler(f)
#        for list in (ul, ur, ml, mr, ll, lr):
#            pickler.dump(list)
#        f.close()
    
    gs = gridspec.GridSpec(3,2)
    gs.update(wspace=0.04, hspace=0.07)
    ax5 = plt.subplot(gs[2, 0])
    ax6 = plt.subplot(gs[2, 1], sharey=ax5)
    ax4 = plt.subplot(gs[1, 1], sharex=ax6)
    ax3 = plt.subplot(gs[1, 0], sharex=ax5, sharey=ax4)
    ax1 = plt.subplot(gs[0, 0], sharex=ax5)
    ax2 = plt.subplot(gs[0, 1], sharex=ax6, sharey=ax1)
    
    
    for ax, data in zip((ax1, ax3, ax5), (ul, ml, ll)):
        ax.hist(data, bins=lbins)
        #AddGamma(ax, data, lrange)
    
    for ax, data in zip((ax2, ax4, ax6), (ur, mr, lr)):
        ax.hist(data, bins=rbins)
        #AddGamma(ax, data, rrange)
    
    ax1.set_ylabel('Low ${\Delta G}_{dsol}$\n(surface) residue count', multialignment = 'center')
    ax3.set_ylabel('Mid ${\Delta G}_{dsol}$\nresidue count', multialignment = 'center')
    ax5.set_ylabel('High ${\Delta G}_{dsol}$\n(buried) residue count', multialignment = 'center')
    ax5.set_xlabel('$\Delta pK_{a}$')
    ax6.set_xlabel('$\Delta\Delta pK_{a}$')
    
    ax5.set_xlim(*lrange)
    ax6.set_xlim(*rrange)
    
    OffTicks(ax1, ['x'])
    OffTicks(ax2, ['x','y'])
    OffTicks(ax3, ['x'])
    OffTicks(ax4, ['x','y'])
    OffTicks(ax5, [])
    OffTicks(ax6, ['y'])
    plt.savefig(os.path.join(outdir, outfile+'.png'))
    
    
#    fig = plt.figure(1)
#    axes = plt.axes()
#
#    axes.set_xlabel('Minimum distance\nto active site $(\AA)$', size = 40, multialignment = 'center')
#    axes.set_ylabel(ylabel, size = 40, multialignment = 'center')
#    fig.set_size_inches(11,8)
#    axes.set_xticklabels(['','10','20','30','40','50'])
#    xtick = axes.get_xticklabels()
#    for tick in xtick:
#        tick.set_fontsize(30)
#    #axes.set_yticklabels([])
#    ytick = axes.get_yticklabels()
#    for tick in ytick:
#        tick.set_fontsize(30)
#    axes.hist(x_high_c, bins=30, color='b')
#    axes.set_xlim(0, 50)
#    fig.savefig(fname+'_highc.png', bbox_inches='tight', transparent=True)
#    axes.cla()
#    
#    axes.set_xlabel('Minimum distance\nto active site $(\AA)$', size = 40, multialignment = 'center')
#    axes.set_ylabel(ylabel, size = 40, multialignment = 'center')
#    fig.set_size_inches(11,8)
#    #axes.set_xticklabels([])
#    #axes.set_yticklabels([])
#    axes.set_xticklabels(['','10','20','30','40','50'])
#    xtick = axes.get_xticklabels()
#    for tick in xtick:
#        tick.set_fontsize(30)
#    axes.set_yticklabels(['0','','1','','2'])
#    ytick = axes.get_yticklabels()
#    for tick in ytick:
#        tick.set_fontsize(30)
#    
#    axes.hist(x_low_c, bins=30, color='r')
#    axes.set_xlim(0, 50)
#    fig.savefig(fname+'_lowc.png', bbox_inches='tight', transparent=True)
#    fig.clf()

if __name__=="__main__":
    dir ='/home/tel/Dropbox/remote/'
    UPSixPanelHisto(dir, 'dpk_ddka_histos')
    ScatterP(dir, 'scatter')
#    data = []
#    for i in range(6):
#        data.append(np.random.rand(10,1)*(np.random.rand()*100))
#    SixPanelHisto(data[0], data[1], data[2], data[3], data[4], data[5])
